package mnet

/*
 * golang task pool.
 */
import (
	"context"
	"fmt"
	"sync"
	"time"
)

const (
	gDefaultWorkerNum    = 8
	gDefaultTaskCapacity = 1024
	gAddTimeoutSecond    = 5
	gExecTimeoutSecond   = 120 // unit:second.
)

// task function: it is a closure function.
type TaskFunc func() interface{}

// TaskResult
type TaskResult struct {
	Result interface{}
	Error  error
}

// TaskWithResult
type TaskWithResult struct {
	Task     TaskFunc
	ResultCh chan TaskResult
	Timeout  time.Duration // 0 is unlimited.
	ID       int64         // task id to assign to the same task channel.
}

// goroutine pool interface.
type IGoroutinePool interface {
	// Start goroutine pool's task.
	Start()

	// Add task with result.
	AddTaskWithResult(task TaskWithResult) error

	// Add task without result and no time limited, the same id will be assigned to the same task queue.
	AddTask(task TaskFunc, ID int64) error

	// Close and wait to close.
	Wait()
}

// goroutinePool defines a goroutine pool.
type goroutinePool struct {
	// Multiple task channels, each channel corresponds to a worker.
	taskChannels []chan TaskWithResult
	workNum      int
	wg           sync.WaitGroup
}

// NewGoroutinePool creates goroutine pool with multiple channels.
func NewGoroutinePool(workerNum int, taskCapacity int) IGoroutinePool {
	if workerNum <= 0 {
		workerNum = gDefaultWorkerNum
	}
	if taskCapacity <= 0 {
		taskCapacity = gDefaultTaskCapacity
	}

	pool := &goroutinePool{
		taskChannels: make([]chan TaskWithResult, workerNum),
		workNum:      workerNum,
	}
	for i := 0; i < workerNum; i++ {
		pool.taskChannels[i] = make(chan TaskWithResult, taskCapacity)
	}

	return pool
}

// Start start to run.
func (p *goroutinePool) Start() {
	for i := 0; i < p.workNum; i++ {
		p.wg.Add(1)

		// Work with i channel.
		go p.worker(i)
	}
}

// getChannelByID get channel with id.
func (p *goroutinePool) getChannelByID(id int64) chan TaskWithResult {
	index := id % int64(p.workNum)
	if index < 0 {
		index = -index
	}
	return p.taskChannels[index]
}

// AddTaskWithResult adds task with result.
func (p *goroutinePool) AddTaskWithResult(task TaskWithResult) error {
	if task.Task == nil {
		return fmt.Errorf("task is nil")
	}

	if task.Timeout < 0 {
		task.Timeout = gExecTimeoutSecond * time.Second
	}

	// get task chan with task id.
	taskChan := p.getChannelByID(task.ID)

	select {
	case taskChan <- task:
		return nil
	case <-time.After(time.Second * gAddTimeoutSecond):
		return fmt.Errorf("add task time out")
	}
}

// AddTask adds task without result.
func (p *goroutinePool) AddTask(task TaskFunc, id int64) error {
	if task == nil {
		return fmt.Errorf("task is nil")
	}

	var taskWithResult TaskWithResult
	taskWithResult.ResultCh = nil // nil is without result.
	taskWithResult.Task = task
	taskWithResult.Timeout = 0 // 0 is unlimited.
	taskWithResult.ID = id
	return p.AddTaskWithResult(taskWithResult)
}

// Wait closes and wait to close.
func (p *goroutinePool) Wait() {
	for _, ch := range p.taskChannels {
		close(ch)
	}
	p.wg.Wait()
}

// worker function.
func (p *goroutinePool) worker(index int) {
	defer p.wg.Done()

	// Deal every task channel of index.
	for task := range p.taskChannels[index] {
		func(task TaskWithResult) {
			defer func() {
				if r := recover(); r != nil {
					if task.ResultCh != nil {
						task.ResultCh <- TaskResult{
							Result: nil,
							Error:  fmt.Errorf("recovered from panic in task: %v", r),
						}
					}
				}
			}()

			// context.
			ctx, cancel := context.WithCancel(context.Background())
			if task.Timeout > 0 {
				ctx, cancel = context.WithTimeout(context.Background(), task.Timeout)
			}
			defer cancel()

			// execute task.
			resultCh := make(chan interface{}, 1)
			errCh := make(chan error, 1)
			go func() {
				defer func() {
					if r := recover(); r != nil {
						errCh <- fmt.Errorf("panic in task: %v", r)
					}
				}()

				result := task.Task()
				select {
				case resultCh <- result:
				case <-ctx.Done():
				}
			}()

			// wait for the result.
			select {
			case result := <-resultCh:
				if task.ResultCh != nil {
					task.ResultCh <- TaskResult{Result: result, Error: nil}
				}
			case err := <-errCh:
				if task.ResultCh != nil {
					task.ResultCh <- TaskResult{Result: nil, Error: err}
				}
			case <-ctx.Done():
				if task.ResultCh != nil {
					task.ResultCh <- TaskResult{Result: nil, Error: ctx.Err()}
				}
			}
		}(task)
	}
}
